#!/usr/bin/env python

# DMigrate - Django Database Migration Tool
# Copyright (C) 2007 Noam Raphael <noamraph at gmail dot com>
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
# 02110-1301, USA.

import sys
import os
import imp
import re
from subprocess import Popen, PIPE
from StringIO import StringIO
from optparse import OptionParser

djangoinst_str = """\
def djangoinst(dir):
    from imp import find_module, load_module
    from django.core.management import setup_environ
    
    settings_mod = load_module('settings', *find_module('settings', [dir]))
    setup_environ(settings_mod)
"""
exec djangoinst_str
    
SEP = "### SEP ###\n\n"

def write_database_settings(out):
    from django.conf import settings
    
    field_names = 'DATABASE_ENGINE DATABASE_HOST DATABASE_NAME '\
                  'DATABASE_PASSWORD DATABASE_PORT DATABASE_USER'.split()
    for field_name in field_names:
        out.write('%s = %r\n' % (field_name, getattr(settings, field_name)))
    out.write('\n')

def get_sqlalchemy_url():
    from django.conf import settings

    engine_dict = {
        'postgresql_psycopg2': 'postgres',
        'postgresql': 'postgres',
        'mysql': 'mysql',
        'mysql_old': 'mysql',
        'sqlite3': 'sqlite',
        'oracle': 'oracle',
        'ado_mssql': 'mssql',
        }

    if settings.DATABASE_ENGINE == 'sqlite3':
        url = 'sqlite:///' + settings.DATABASE_NAME
    else:
        if settings.DATABASE_HOST == '':
            host = 'localhost'
        elif settings.DATABASE_HOST.startswith('/'):
            raise ValueError, 'Unix sockets not supported as DATABASE_HOST'
        else:
            host = settings.DATABASE_HOST
        url = '%(engine)s://%(user)s:%(password)s@%(host)s%(port)s/%(name)s' \
              % dict(engine = engine_dict[settings.DATABASE_ENGINE],
                     user = settings.DATABASE_USER,
                     password = settings.DATABASE_PASSWORD,
                     host = host,
                     port = ':'+settings.DATABASE_PORT if settings.DATABASE_PORT
                            else '',
                     name = settings.DATABASE_NAME,
                     )
    return url
        

def write_field_list(out, apps_models):
    for app, models in apps_models:
        out.write("%s\n" % get_app_label(app))
        for model in models:
            opts = model._meta
            out.write("    %s\n" % opts.object_name)
            for field in opts.fields + opts.many_to_many:
                out.write("        %s\n" % field.name)
    out.write('\n')

def parse_field_list(field_list):
    apps = []
    models = None
    fields = None

    for line in field_list.split('\n'):
        if not line.strip():
            continue
        nspaces = len(re.match(r'^ *', line).group())
        name = line.strip()
        if nspaces == 0:
            fields = None
            models = []
            apps.append((name, models))
        elif nspaces == 4:
            if models is None:
                raise ValueError, \
                      "Error parsing field list - model without an app"
            fields = []
            models.append((name, fields))
        elif nspaces == 8:
            if fields is None:
                raise ValueError, \
                      "Error parsing field list - field without a model"
            fields.append(name)
        else:
            raise ValueError, \
                  "Error parsing field list - wrong number of spaces"

    return apps


data_types = {'AutoField': 'Integer',
              'BooleanField': 'Boolean',
              'CharField': 'VARCHAR(%(max_length)s)',
              'CommaSeparatedIntegerField': 'VARCHAR(%(max_length)s)',
              'DateField': 'Date',
              'DateTimeField': 'DateTime',
              'DecimalField': 'DECIMAL',
              'FileField': 'VARCHAR(100)',
              'FilePathField': 'VARCHAR(100)',
              'FloatField': 'Float',
              'IPAddressField': 'CHAR(15)',
              'ImageField': 'VARCHAR(100)',
              'IntegerField': 'Integer',
              'ManyToManyField': None,
              'NullBooleanField': 'Boolean',
              'OneToOneField': 'Integer',
              'PhoneNumberField': 'VARCHAR(20)',
              'PositiveIntegerField': 'Integer',
              'PositiveSmallIntegerField': 'SmallInteger',
              'SlugField': 'VARCHAR(%(max_length)s)',
              'SmallIntegerField': 'SmallInteger',
              'TextField': 'TEXT',
              'TimeField': 'Time',
              'USStateField': 'VARCHAR(2)'}
def get_data_type(f):
    internal_type = f.get_internal_type()
    return data_types[internal_type] % f.__dict__

def get_table_obj_name(db_table):
    return db_table+'_table'

def get_src_class_obj_name(opts):
    return 'src.%s.%s' % (opts.app_label, opts.object_name)

def get_dst_class_obj_name(opts):
    return 'dst.%s.%s' % (opts.app_label, opts.object_name)

def write_sqlalchemy_tables(out, model):
    from django.db import models as M
    from django.contrib.contenttypes import generic

    opts = model._meta

    out.write("%s = Table(%r, metadata,\n" %
              (get_table_obj_name(opts.db_table), opts.db_table))

    for f in opts.fields:
        if isinstance(f, M.ManyToManyField):
            continue
        
        if isinstance(f, (M.ForeignKey, M.OneToOneField)):
            rel_field = f.rel.get_related_field()
            while isinstance(rel_field, (M.ForeignKey, M.OneToOneField)):
                rel_field = rel_field.rel.get_related_field()
        else:
            rel_field = f

        out.write("    Column(%r, %s" % (f.column, get_data_type(rel_field)))
        if f.rel:
            out.write(", ForeignKey(%r)" %
                      (f.rel.to._meta.db_table + '.' +
                       f.rel.to._meta.get_field(f.rel.field_name).column))
        if f.primary_key:
            out.write(", primary_key=True")
        if not f.null and not f.primary_key:
            out.write(", nullable=False")
        if f.unique:
            out.write(", unique=True")
        out.write("),\n")

    for field_constraints in opts.unique_together:
        out.write("\n")
        out.write("    UniqueConstraint(%s),\n" %
                  ', '.join(repr(opts.get_field(f).column)
                            for f in field_constraints))

    out.write("    )\n\n")

    for f in opts.many_to_many:
        if not isinstance(f.rel, generic.GenericRel):
            out.write("""\
%s = Table(%r, metadata,
    Column('id', Integer, primary_key=True),
    Column(%r, %s, ForeignKey(%r), nullable=False),
    Column(%r, %s, ForeignKey(%r), nullable=False),
    UniqueConstraint(%r, %r),
    )
    
""" %
                      (get_table_obj_name(f.m2m_db_table()),
                       f.m2m_db_table(),
                       
                       f.m2m_column_name(),
                       get_data_type(opts.pk),
                       opts.db_table+'.'+opts.pk.column,

                       f.m2m_reverse_name(),
                       get_data_type(f.rel.to._meta.pk),
                       f.rel.to._meta.db_table+'.'+f.rel.to._meta.pk.column,

                       f.m2m_column_name(),
                       f.m2m_reverse_name(),
                       ))


def get_app_label(app):
    """Given an app (the models module), return the app_label."""
    return app.__name__.split('.')[-2]

def write_sqlalchemy_classes(out, apps_models):
    out.write('class src(object):\n\n')
    for app, models in apps_models:
        out.write('    class %s(object):\n' % get_app_label(app))
        for model in models:
            out.write('        class %s(object):\n' % model._meta.object_name)
            out.write('            pass\n')
        out.write('        \n')
    out.write('    \n')


def write_sqlalchemy_mappers(out, model):
    from django.db import models as M
    from django.contrib.contenttypes import generic

    opts = model._meta
    
    out.write("mapper(%s, %s, properties=dict(\n" %
              (get_src_class_obj_name(opts),
               get_table_obj_name(opts.db_table),
               ))
    
    for f in opts.fields:
        if not isinstance(f, M.ManyToManyField) and f.rel:
            backref = f.rel.related_name or (opts.object_name.lower()+'_set')
            join_condition = '%s.c.%s==%s.c.%s' % (
                get_table_obj_name(f.rel.to._meta.db_table),
                f.rel.field_name,
                get_table_obj_name(opts.db_table),
                f.column)
            out.write("    %s = relation(%s, backref=%r,\n"
                      "        primaryjoin=%s),\n" %
                      (f.name,
                       get_src_class_obj_name(f.rel.to._meta),
                       backref,
                       join_condition))

    for f in opts.many_to_many:
        if not isinstance(f.rel, generic.GenericRel):
            backref = f.rel.related_name or (opts.object_name.lower()+'_set')
            out.write("    %s = relation(%s, backref=%r,\n"
                      "        secondary=%s),\n" %
                      (f.name,
                       get_src_class_obj_name(f.rel.to._meta),
                       backref,
                       get_table_obj_name(f.m2m_db_table()),
                       ))

    out.write("    ))\n\n")
    

def write_django_classes(out, apps_models):
    out.write('class dst(object):\n\n')
    for app, models in apps_models:
        out.write('    class %s(object):\n' % get_app_label(app))
        out.write('        from %s import %s\n' %
                  (app.__name__,
                   ', '.join(model._meta.object_name for model in models)))
        out.write('        \n')
    out.write('    \n')


def write_django_fix_auto(out, apps_models):
    from django.db.models.fields import DateField, TimeField
    for app, models in apps_models:
        for model in models:
            opts = model._meta
            for field in opts.fields:
                if isinstance(field, (DateField, TimeField)):
                    if field.auto_now:
                        out.write('%s._meta.get_field(%r).auto_now = False\n'
                                  % (get_dst_class_obj_name(opts),
                                     field.name))
                    if field.auto_now_add:
                        out.write('%s._meta.get_field(%r).auto_now_add = False\n'
                                  % (get_dst_class_obj_name(opts),
                                     field.name))
    out.write('\n')


def parse_renames(renames):
    """
    Get a rename string, in the format
    app.SrcClass-app.DstClass,app.SrcClass-app.DstClass

    and return a list of tuples of tuples of strings:
    [(('app', 'SrcClass'), ('app', 'DstClass')),
     (('app', 'SrcClass'), ('app', 'DstClass')),
     ]
    """
    r = []
    for rename in renames.split(','):
        sp = rename.split('-')
        if len(sp) != 2:
            raise ValueError, "Rename unit %r didn't contain one '-'" % rename
        src, dst = sp

        sp = src.split('.')
        if len(sp) != 2:
            raise ValueError, "Source model %r didn't contain one '.'" % src
        src_app, src_model = sp

        sp = dst.split('.')
        if len(sp) != 2:
            raise ValueError, "Dest model %r didn't contain one '.'" % dst
        dst_app, dst_model = sp

        r.append(((src_app.strip(), src_model.strip()),
                  (dst_app.strip(), dst_model.strip())))
    return r
        

def get_modifications(renames, field_list, apps_models):
    """
    Return a tuple like this:
    ([removed models],
     [added models],
     {mapping between models which stay, to:
      (old model,
       [removed field names],
       [added field names],
       [field names which stay]
       )
      }
     )

    A "model", in this context, is a tuple (app_label, model_name).
    """
    renames = dict(renames)
    renames_values = set(renames.itervalues())
    rename_count = 0
    dst_models = set((get_app_label(app), model._meta.object_name)
                     for app, models in apps_models
                     for model in models)

    removed_models = []

    src_models = {}
    for app_label, models in field_list:
        for model_name, fields in models:
            orig_name = (app_label, model_name)
            if orig_name in renames:
                new_name = renames[orig_name]
                rename_count += 1
                if new_name not in dst_models:
                    raise ValueError, "Dest model %s.%s isn't found" % new_name
            else:
                if orig_name not in renames_values:
                    new_name = orig_name
                else:
                    new_name = None
                if new_name not in dst_models:
                    removed_models.append(orig_name)
                    new_name = None
            if new_name:
                src_models[new_name] = (orig_name, fields)

    if rename_count != len(renames):
        raise ValueError, "Source model %s.%s isn't found" % \
              renames.popitem()[0]

    added_models = []
    stay_models = {}

    for app, models in apps_models:
        app_label = get_app_label(app)
        for model in models:
            opts = model._meta
            tup = (app_label, opts.object_name)
            if tup not in src_models:
                added_models.append(tup)
            else:
                orig_name, src_fields = src_models[tup]
                dst_fields = [x.name for x in opts.fields + opts.many_to_many]
                removed_fields = [x for x in src_fields
                                  if x not in dst_fields]
                added_fields = [x for x in dst_fields
                                if x not in src_fields]
                stay_fields = [x for x in dst_fields
                               if x in src_fields]
                stay_models[tup] = (orig_name,
                                    removed_fields, added_fields, stay_fields)

    return removed_models, added_models, stay_models
    

def write_convert(out, modifications, apps_models):
    removed_models, added_models, stay_models = modifications

    from django.db.models import ForeignKey, ManyToManyField
    
    app_labels = [get_app_label(app) for app, models in apps_models]
    out.write('''\
def convert():
    """Reset the destination DB, and copy all data from the source DB to it,
    with the necessary changes.

    The user should change this function so that it will make the changes.
    Follow the TO-DO comments! (Without the '-'...)
    """
    # Reset the destination DB
    
    r = raw_input("Are you sure you want to reset the destination database? ")
    if r.lower() not in ('y', 'yes'):
        print >> sys.stderr, "Quitting."
        sys.exit(1)
    from django.db.models import get_app
    from django.core.management import load_command_class
    reset = load_command_class('django.core', 'reset').handle_app
    for app_label in %r:
        reset(get_app(app_label), interactive=False)
    
''' % app_labels)

    if removed_models or added_models:
        out.write("    # TODO:\n")
        if removed_models:
            out.write("    # Removed models: %s\n" %
                      ', '.join('src.%s.%s' % model
                                for model in removed_models))
        if added_models:
            out.write("    # Added models: %s\n" %
                      ', '.join('dst.%s.%s' % model
                                for model in added_models))
        if removed_models and added_models:
            out.write("    # If you renamed a model, use the --rename option!\n")
        out.write("    \n")

    out.write('''\
    # Copy data!
    
''')

    for app, models in apps_models:
        app_label = get_app_label(app)
        for model in models:
            opts = model._meta
            try:
                data = stay_models[app_label, opts.object_name]
            except KeyError:
                # An added model, do nothing
                continue
            orig_name, removed_fields, added_fields, stay_fields = data
                       
            out.write("    for s in progress(session.query(src.%s.%s)):\n"
                      % orig_name)
            out.write("        d = %s()\n" % get_dst_class_obj_name(opts))

            for field_name in stay_fields:
                field = opts.get_field(field_name)

                if isinstance(field, ForeignKey):
                    out.write("        d.%s = s.%s\n" %
                              (field.get_attname(), field.get_attname()))
                elif isinstance(field, ManyToManyField):
                    out.write("        for x in s.%s:\n" % field.name)
                    out.write("            d.%s.add(x.%s)\n" %
                              (field.name, field.rel.to._meta.pk.attname))
                else:
                    out.write("        d.%s = s.%s\n" %
                              (field.name, field.name))
            out.write("        \n")

            if removed_fields or added_fields:
                out.write("        # TODO:\n")
                if removed_fields:
                    out.write("        # Removed fields: %s\n" %
                              ', '.join(removed_fields))
                if added_fields:
                    out.write("        # Added fields: %s\n" %
                              ', '.join(added_fields))
                out.write("        \n")

            out.write("        d.save()\n")
            
            out.write("    \n")
            

def get_apps_models():
    from django.db.models.loading import get_apps, get_models

    r = [(app, get_models(app)) for app in get_apps()]

    if len(set(get_app_label(app) for app, models in r)) != len(r):
        raise NotImplementedError, "Can't have two apps with the same name"
    
    return r


def main_src():
    djangoinst(sys.argv[2])
    
    from django.core.management.validation import get_validation_errors

    out = sys.stdout

    print >> out, get_sqlalchemy_url()
    print >> out
    out.write(SEP)

    # Validate current models
    outfile = StringIO()
    if get_validation_errors(outfile) != 0:
        print >> sys.stderr, "There are validation errors in the source project."
        print >> sys.stderr, "Run './manage.py validate' to see them."
        sys.exit(1)

    apps_models = get_apps_models()

    write_field_list(out, apps_models)
    out.write(SEP)

    print >> out, "# Define src database tables"
    print >> out

    for app, models in apps_models:
        for model in models:
            write_sqlalchemy_tables(out, model)

    print >> out, "# Define src database classes"
    print >> out

    write_sqlalchemy_classes(out, apps_models)

    print >> out, "# Define src database mappers"
    print >> out

    for app, models in apps_models:
        for model in models:
            write_sqlalchemy_mappers(out, model)


def main():
    # Run source process, if first arg is "_src".
    if len(sys.argv) > 1 and sys.argv[1] == '_src':
        return main_src()

    # Parse args

    usage = '%prog [options] [src_dir] dst_dir migrate_script'
    parser = OptionParser(usage=usage)
    parser.add_option('-r', '--rename', dest='renames',
                      help='List of model renames, in the format '
                      'app.SrcClass-app.DstClass,app.SrcClass-app.DstClass')
    parser.add_option('-f', '--force', dest='force', action='store_true',
                      default=False,
                      help="Don't ask whether to overwrite migrate_script.")
    options, args = parser.parse_args()

    if len(args) not in (2, 3):
        parser.error('2 or 3 arguments are expected')
    if len(args) == 2:
        dst_dir, migrate_script = args
        src_dir = dst_dir
    else:
        src_dir, dst_dir, migrate_script = args

    if options.renames:
        renames = parse_renames(options.renames)
    else:
        renames = []

    if not options.force and os.path.exists(migrate_script):
        r = raw_input("%s already exists. Overwrite? " % migrate_script)
        if r.lower() not in ('y', 'yes'):
            print >> sys.stderr, "Quitting."
            sys.exit(1)

    # Run source subprocess and get its output

    project_name = os.path.basename(os.path.abspath(dst_dir))
    p = Popen([sys.executable, __file__, '_src', src_dir, project_name],
              stdout=PIPE)
    inp = p.stdout.read()
    rc = p.wait()
    if rc != 0:
        print >> sys.stderr, "Source subprocess returned error code. Exiting"
        sys.exit(rc)

    # Configure django to dest project.

    djangoinst(dst_dir)
    from django.conf import settings
    from django.core.management.validation import get_validation_errors

    splitted_inp = inp.split(SEP)
    if len(splitted_inp) != 3:
        print >> sys.stderr, "Subprocess output wasn't in expected format"
        sys.exit(1)

    sa_url, field_list, sa_code = splitted_inp
    sa_url = sa_url.strip()
    field_list = parse_field_list(field_list)

    # Validate current models
    outfile = StringIO()
    if get_validation_errors(outfile) != 0:
        print >> sys.stderr, "There are validation errors in the dest project."
        print >> sys.stderr, "Run './manage.py validate' to see them."
        sys.exit(1)

    apps_models = get_apps_models()

    out = open(migrate_script, 'w')

    print >> out, '''\
#!%s
# -*- coding: utf-8 -*-

# This is a Django migration script, automatically created by DMigrate.
#
# Before running it, search for TODO comments, which mark new fields and
# deleted fields.
# You should also check the source database settings and the destination
# database settings.

%s
djangoinst(%r)

# Configure source database

SRC_DB_URL = %r

# Configure dest database. If you want to override the settings in the
# settings file, uncomment these lines:

from django.conf import settings
dummy = settings.DATABASE_NAME # Work around bug #4884
''' % \
    (sys.executable, djangoinst_str, os.path.abspath(dst_dir), sa_url)
    
    field_names = 'DATABASE_ENGINE DATABASE_HOST DATABASE_NAME '\
                  'DATABASE_PASSWORD DATABASE_PORT DATABASE_USER'.split()
    for field_name in field_names:
        print >> out, '# settings.%s = %r' % (field_name, getattr(settings, field_name))
        
    print >> out, '''\

# import statements

import sys

from sqlalchemy import *
from sqlalchemy.orm import sessionmaker, mapper, relation


# Tell SQLAlchemy how to access the source database

engine = create_engine(SRC_DB_URL)
Session = sessionmaker(bind=engine, autoflush=True, transactional=True)
session = Session()
metadata = MetaData()
'''
    print >> out, sa_code

    print >> out, '# Import dst database classes'
    print >> out

    write_django_classes(out, apps_models)

    print >> out, '# Temporarily disable auto_now and auto_now_add'
    print >> out

    write_django_fix_auto(out, apps_models)

    print >> out, r"""
def progress(query):
    count = query.count()
    table_name = query.table.name
    max_str_len = 0
    for i, x in enumerate(query):
        s = "%s %d/%d" % (table_name, i, count)
        max_str_len = max(max_str_len, len(s))
        sys.stderr.write(s+" "*(max_str_len-len(s))+"\r")
        yield x
    sys.stderr.write(" "*max_str_len + "\r")
"""
    print >> out, '# The great conversion function'
    print >> out

    modifications = get_modifications(renames, field_list, apps_models)
    write_convert(out, modifications, apps_models)
    
    print >> out, '''\
def main():
    if len(sys.argv) != 1:
        print >> sys.stderr, """\\
Usage: %s
Migrate a Django database.
See the source - You'll most likely want to edit it!""" % sys.argv[0]
        sys.exit(1)
    else:
        convert()

if __name__ == '__main__':
    main()
'''

    out.close()

    # Give execute permissions to everyone who has read permissions
    cur_mode = os.stat(migrate_script).st_mode
    new_mode = cur_mode | ((cur_mode & 0444) >> 2)
    os.chmod(migrate_script, new_mode)
    

if __name__ == '__main__':
    main()

